import argparse
import os
import onnxruntime
from PIL import Image  
import numpy as np
import shutil

def create_parser():
  parser = argparse.ArgumentParser(description='Image Classification using ResNet ONNX Model')
  parser.add_argument('--model', type=str, required=True, help='Path to the ONNX model file')
  parser.add_argument('--class-file', type=str, required=True, help='Path to the class names file')
  parser.add_argument('--input-folder', type=str, required=True, help='Path to the input image folder')
  parser.add_argument('--output-folder', type=str, required=True, help='Path to the output folder')
  parser.add_argument('--img-size', type=int, default=224, help='Input image size (default: 224)')
  parser.add_argument('--mean', type=float, nargs='+', default=[0.485, 0.456, 0.406], help='Mean pixel value for preprocessing (default: [0.485, 0.456, 0.406])')
  parser.add_argument('--std', type=float, nargs='+', default=[0.229, 0.224, 0.225], help='Standard deviation for preprocessing (default: [0.229, 0.224, 0.225])')
  parser.add_argument('--gpu', action='store_true', default=False, help='Use GPU for inference')
  return parser

def main(args):
  
  # Set ONNX Runtime execution provider
  if args.gpu:
    providers = ['CUDAExecutionProvider']
  else:
    providers = ['CPUExecutionProvider']

  # Load ONNX model
  session = onnxruntime.InferenceSession(args.model, providers=providers)

  # Load class names
  class_names = []
  with open(args.class_file, 'r') as f:
    class_names = f.read().splitlines()

  # Create output folders
  for name in class_names:
    folder = os.path.join(args.output_folder, name)
    os.makedirs(folder, exist_ok=True)

  # Classify images
  for root, _, files in os.walk(args.input_folder):
    for file in files:
      if file.endswith(".jpg"):
      
        image_path = os.path.join(root, file)
        
        # Load image
        img = Image.open(image_path).convert('RGB')
        
        # Preprocess image
        img = img.resize((args.img_size, args.img_size), Image.BILINEAR)
        img = np.array(img, dtype=np.float32) / 255 # Use float32
        img = (img - args.mean) / args.std
        img = np.transpose(img, (2, 0, 1))
        img = np.expand_dims(img, 0)

        # Cast to float32 right before inference 
        img = img.astype(np.float32) 

        # Run inference
        result = session.run(None, {'input0': img})
        predicted_class = np.argmax(result)
        
        # Get class name
        class_name = class_names[predicted_class]
        
        # Copy image to output folder
        output_path = os.path.join(args.output_folder, class_name, os.path.basename(image_path))  
        shutil.copy(image_path, output_path)

        print(f"Image '{file}' classified as '{class_name}' and copied to '{class_name}' folder.")

  print("Classification and copying completed.")

if __name__ == '__main__':
  parser = create_parser()
  args = parser.parse_args()
  main(args)